#include <iomanip>
#include <exception>
#include <cmath>
#include "Core/Wrapper_LAPACKE.h"
#include "FiniteDiff/MixedDer.h"
#include "Embedded/EMixedDer.h"
#include "Embedded/SourceStencil.impl.h"
#include "Embedded/EvaluateMonomials.h"
#include "Embedded/ConditionNumber.h"
#include "../example/common/TestUtility.h"

#if DIM == 2

struct EMixedDer::MixedDerTransform
{
  Vec<Real, 3> A;
  MixedDerTransform(const Vec<Real, 3> &aA) : A(aA)
  {
  }

  bool checkStandardScheme(const Tensor<int, Dim> &nodeLabels, Vec<int, Dim> idx, bool useOuterRect) const
  {
    auto validBox = nodeLabels.box();
    if(idx[0] < validBox.lo()[0] + 2 || idx[0] > validBox.hi()[0] - 2
        || idx[1] < validBox.lo()[1] + 2 || idx[1] > validBox.hi()[1] - 2)
      return useOuterRect;
//    using NodeType = typename EmbeddedGrid<Dim>::NodeType;
    for(int j = -2; j <= 2; ++j)
      for(int i = -2; i <= 2; ++i)
        if(nodeLabels(idx + iVec {i, j}) < 0)
          return false;
    return true;
  }

  template <class... Ts>
  Real transform(const Vec<Real, Dim> &p, Ts... indices) const {
    Vec<int, Dim> beta { decltype(indices)::value ... };
    Real xx, xy, yy;
    // xx term
    if (beta[0] < 2) {
      xx = 0.0;
    } else {
      beta[0] -= 2;
      xx = 1.0;
      for (int d = 0; d < Dim; ++d)
        xx *= std::pow(p[d], beta[d]);
      beta[0] += 2;
      xx *= beta[0] * (beta[0]-1);
    }
    // xy term
    if(beta[0] < 1 || beta[1] < 1) {
      xy = 0.0;
    } else {
      --beta[0];
      --beta[1];
      xy = 1.0;
      for(int d = 0; d < Dim; ++d)
        xy *= std::pow(p[d], beta[d]);
      ++beta[0];
      ++beta[1];
      xy *= beta[0] * beta[1];
    }
    // yy term
    if(beta[1] < 2) {
      yy = 0.0;
    } else {
      beta[1] -= 2;
      yy = 1.0;
      for(int d = 0; d < Dim; ++d)
        yy *= std::pow(p[d], beta[d]);
      beta[1] += 2;
      yy *= beta[1] * (beta[1]-1);
    }
    //
    return A[0]*xx + A[1]*xy + A[2]*yy;
  }
};

//======================================================================

EMixedDer::EMixedDer(const EmGrid &aGrid, const Vec<Real, 3> &aA)
    : BaseClass(aGrid, aA), emGrid(aGrid)
{
  // 0.
  struct _MDTransform : public MixedDerTransform {
    enum { fitOrder = EMixedDer::fitOrder };
    _MDTransform(const Vec<Real, 3> &aA) : MixedDerTransform(aA) { }
  };
  _MDTransform _mdt(aA);
  // 1. Construct the stencils.
  SourceStencilParam<Dim> ssparam;
  ssparam.otype = SourceStencilParam<Dim>::OperatorType::Elliptic;
  CPUTimer timer;
  timer.reset();
  SourceStencil<Dim> ss(emGrid, _mdt, ssparam);
  std::cout << "Construction of stencils : " << std::setprecision(5) << timer() << "s. " << std::endl << std::defaultfloat;
#ifndef NDEBUG
  if(get_dbglevel() >= 3)
    ss.dump(tmpos);
#endif

  // 2. Find the approximation.
  timer.reset();
//  using NodeType = typename EmGrid::NodeType;
  const auto &bdryPoints = emGrid.bdryPoints;
//  const auto &outNormals = emGrid.outNormals;
  const auto dx = emGrid.spacing();
  const int n = fitOrder - 1;
  const int nBasis = binom(n+Dim, Dim);
  auto DirichletModifier = [&](const Real &x, const auto &) { return x; };

  for(const auto &stencil : ss.vStencil) {
    const int nSources = stencil.locations.size();
    const auto destIdx = stencil.destIdx;
    const auto refLoc  = (destIdx + 0.5) * dx;
    // 2.1. calculate the sample matrix.
    Tensor<Real, 2> sampleMatrix(Vec<int, 2> {nBasis, nSources});
    sampleMatrix = 0.0;
    Real *where2Fill = sampleMatrix.data();
    for(std::size_t k = 0; k < stencil.locations.size(); ++k) {
      rVec p;
      if(stencil.isBdry[k] == 0) {
        iVec jdx = stencil.locations[k];
        p = (jdx + 0.5) * dx;
        EvaluateMonomials<Dim,n>::doEvaluate(identityOfMonomials(), DirichletModifier, where2Fill, p-refLoc, rVec(0.0), dx);
      } else {
        // use the nearest boudary point as the boundary condition
        p = bdryPoints(destIdx);
        EvaluateMonomials<Dim,n>::doEvaluate(identityOfMonomials(), DirichletModifier, where2Fill, p-refLoc, rVec(0.0), dx);
      }
    }
    assert((where2Fill - sampleMatrix.data()) == sampleMatrix.volume());
#ifndef NDEBUG
    const Real cn = estimateUnderdetCond(sampleMatrix);
    if(cn >= dangerousCN || get_dbglevel() >= 2)
      dbgcout1 << "The linear system for " << destIdx
               << " has condition number " << std::setprecision(2) << std::scientific << cn
               << ". \n" << std::defaultfloat;
    dbgcout3 << std::showpos << std::setprecision(3) << sampleMatrix << "\n";
    dbgcout3 << std::defaultfloat;
#endif
    // 2.2. calculate the RHS.
    Tensor<Real, 1> sampleRHS(nSources);
    sampleRHS = 0.0;
    where2Fill = sampleRHS.data();
    EvaluateMonomials<Dim,n>::doEvaluate(_mdt, DirichletModifier, where2Fill, rVec(0.0), rVec(0.0), dx);
    // 2.3. solve the linear system to obtain the coefficients.
    dbgcout3 << sampleMatrix << std::endl;
    // scale the columns
    for(int j = 0; j < nSources; ++j)
      for(int i = 0; i < nBasis; ++i)
        sampleMatrix(i, j) /= stencil.weights[j];
    auto info = LAPACKE_dgels(LAPACK_COL_MAJOR, 'N', nBasis, nSources, 1, sampleMatrix.data(), nBasis, sampleRHS.data(), nSources);
    if(info != 0)
      throw std::runtime_error("(PoissonOp::PoissonOp) DGELS failed.");
    // scale the raw solution
    for(int j = 0; j < nSources; ++j)
      sampleRHS(j) /= stencil.weights[j];
    // 2.4. Save the coefficients in sparse matrix format.
    matOp.insertRow(destIdx, nSources, stencil.locations.data(), sampleRHS.data());
    dbgcout2 << sampleRHS << "\n";
  }
  std::cout << "Discretization of the operator : " << std::setprecision(5) << timer() << "s. " << std::endl << std::defaultfloat;
}

void EMixedDer::apply(const Tensor<Real, Dim> &aData, Tensor<Real, Dim> &outData) const
{
  BaseClass::apply(aData, outData);
  matOp.AXPY(1.0, aData, [](iVec) { return 0.0; }, outData);
}

void EMixedDer::apply(const Tensor<Real, Dim> &aData, const Tensor<Real, Dim> &rhs, Tensor<Real, Dim> &rsd) const
{
  BaseClass::apply(aData, rhs, rsd);
  matOp.AXPY(-1.0, aData, rhs, rsd);
}

#endif // DIM == 2
